package com.pshdhx.aiagent.chatMemory;

import com.esotericsoftware.kryo.Kryo;
import com.esotericsoftware.kryo.io.Input;
import com.esotericsoftware.kryo.io.Output;
import lombok.extern.slf4j.Slf4j;
import org.objenesis.strategy.InstantiatorStrategy;
import org.objenesis.strategy.StdInstantiatorStrategy;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;
import org.springframework.stereotype.Component;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;

/**
 * @Author pansd
 * @Date 2025-06-27 17:09
 * @Des
 */
@Component
@Slf4j
public class SaveChatBySerializable implements ChatMemory {

    private static final Kryo kryo = new Kryo();
    private static final String BASE_DIR = System.getProperty("user.dir") + "/chatMemory";

    static {
        kryo.setRegistrationRequired(false);
        // 设置实例化策略
        kryo.setInstantiatorStrategy(new StdInstantiatorStrategy());
    }

    public SaveChatBySerializable() {
        File file = new File(BASE_DIR);
        if (!file.exists()) {
            file.mkdirs();
        }
    }

//    @Override
//    public void add(String conversationId, Message message) {
//        try {
//            serializeToFile(List.of(message), conversationId);
//        } catch (IOException e) {
//            e.printStackTrace();
//        }
//    }

    @Override
    public void add(String conversationId, List<Message> messages) {
        try {
            List<Message> messageList = deserializeFromFile(conversationId);
            messageList.addAll(messages);
            serializeToFile(messageList, conversationId);
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    @Override
    public List<Message> get(String conversationId, int lastN) {
        try {
            List<Message> messages = deserializeFromFile(conversationId);
//            return messages.stream().skip(messages.size() - lastN).collect(Collectors.toList());

            //此处判断有问题
            return messages.stream()
                    .skip(Math.max(0, messages.size() - lastN))
                    .toList();
        } catch (IOException e) {
            e.printStackTrace();
        }
        return null;
    }

    @Override
    public void clear(String conversationId) {
        File file = new File(BASE_DIR, conversationId + ".kryo");
        if (file.exists()) {
            file.delete();
        }
    }

    /**
     * 序列化 List<Message> 到文件
     */
    public static void serializeToFile(List<Message> messages, String filename) throws IOException {
        try (Output output = new Output(new FileOutputStream(new File(BASE_DIR, filename + ".kryo")))) {
            kryo.writeObject(output, messages);
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    /**
     * 从文件反序列化为 List<Message>
     * 读取时，缺少了判断文件是否存在
     */
    public static List<Message> deserializeFromFile(String filename) throws IOException {
        List<Message> messages = new ArrayList<>();
        File file = new File(BASE_DIR, filename + ".kryo");
        if (file.exists()) {
            try (Input input = new Input(new FileInputStream(file))) {
                messages = kryo.readObject(input, ArrayList.class);
            }
        }
        return messages;
    }

}
